//===------------------- Damsl Lab @ UD CIS -------------------------------===//
//
//                     MPI Optimizaition Framework
//
//
//
//
//===----------------------------------------------------------------------===//
//
// This pass implements preliminary transformation in Anthony's paper
//
//
//===----------------------------------------------------------------------===//


#define DEBUG_TYPE "mpi-pre-transform"
#include "llvm/Transforms/Scalar.h"
#include "llvm/Analysis/LoopPass.h"
#include "llvm/Transforms/Utils/UnrollLoop.h"
#include "llvm/Support/Debug.h"
#include "llvm/Analysis/DebugInfo.h"
#include "llvm/Transforms/MPI/MPI.h"
#include "llvm/Metadata.h"
#include <vector>
#include <algorithm>

using namespace llvm;

namespace {
class MPIPreTransform: public LoopPass {
	bool isMpiCommCall(const Instruction &II);
	bool isMpiCommOnlyLoop(const Loop *L);
public:
	static char ID;
	MPIPreTransform() : LoopPass(ID) {}
	bool runOnLoop(Loop *L, LPPassManager &LPM);
	virtual void getAnalysisUsage(AnalysisUsage &AU) const {
		AU.addRequiredID(LoopSimplifyID);
		AU.addRequiredID(LCSSAID);
		AU.addPreservedID(LCSSAID);
		AU.addRequired<LoopInfo> ();
		AU.addPreserved<LoopInfo> ();
		// FIXME: Loop unroll requires LCSSA. And LCSSA requires dom info.
		// If loop unroll does not preserve dom info then LCSSA pass on next
		// loop will receive invalid dom info.
		// For now, recreate dom info, if loop is unrolled.
		AU.addPreserved<DominatorTree> ();
		AU.addPreserved<DominanceFrontier> ();
	}
};
}

char MPIPreTransform::ID = 0;
static RegisterPass<MPIPreTransform> X("mpi-pre-transform", "MPI Preliminary Transformation Pass");
//INITIALIZE_PASS(MPIPreTransform, "mpi-pre-transform", "MPI Preliminary Transformation Pass", false, false);

// Return true if Instruction II is a MPI Communication call.
bool MPIPreTransform::isMpiCommCall(const Instruction &II) {
	if (const CallInst* callInst = dyn_cast<CallInst>(&II)) {
		Function* func = callInst->getCalledFunction();
		StringRef fname = func->getName();
		if (fname.startswith(StringRef("MPI_Send"))	 ||
			fname.startswith(StringRef("MPI_Isend")) ||
			fname.startswith(StringRef("MPI_Recv"))  ||
			fname.startswith(StringRef("MPI_Irecv")))
			return true;
	}
	return false;
}

bool MPIPreTransform::isMpiCommOnlyLoop(const Loop *L) {
//	std::vector<BasicBlock*> BV = L->getBlocks();
//	for (std::vector<BasicBlock*>::iterator BI = BV.begin(), BE = BV.end(); BI
//			!= BE; ++BI) {
//		for (BasicBlock::iterator II = (*BI)->begin(), E = (*BI)->end(); II
//				!= E; ++II) {
//			if (this->isMpiCommCall(*II)) {
//				if(MDNode *N = II->getMetadata("whether-mpi-only-loop")) {
//					Value *t = N->getOperand(0);
//					ConstantInt * a1 = dyn_cast<ConstantInt>(t);
//					ConstantInt * a2 = ConstantInt::getTrue((*BI)->getContext());
//					if(a1 == a2) return true;
//				}
//			}
//		}
//	}
//	return false;


 return true;

}

bool MPIPreTransform::runOnLoop(Loop *L, LPPassManager &LPM) {
	LoopInfo *LI = &getAnalysis<LoopInfo> ();

	if (L->getSubLoops().size() != 0)
		return false;

	// TODO: getSmallConstantTripCount requires induction variable is in
	// canonical form, And the first instruction in loop basicblock in a phi
	// instruction. Details in /lib/Analysis/LoopInfo.cpp.
	// To make getSmallConstantTripCount working, O1 optimization is required
	// beforehand.
	unsigned TripCount = L->getSmallConstantTripCount();
	unsigned Count = 0;
	DEBUG(errs() << "TripCount is " << TripCount << "\n");

	// Conservative heuristic: if we know the trip count, see if we can
	// completely unroll.
	// make sure that this loop have statically known iteration space.
	// Currently, the iteration space is one dimension
	if (TripCount == 0)
		return false;
	Count = TripCount;

	if (!isMpiCommOnlyLoop(L)) return false;

	// Unroll the loop.
	Function *F = L->getHeader()->getParent();
	if (!UnrollLoop(L, Count, LI, &LPM)) {
		DEBUG(errs() << "Unroll loop failed.\n");
		return false;
	}
	DEBUG(errs() << "Unroll loop success.\n");

	// FIXME: Reconstruct dom info, because it is not preserved properly.
	DominatorTree *DT = getAnalysisIfAvailable<DominatorTree> ();
	if (DT) {
		DT->runOnFunction(*F);
		DominanceFrontier *DF = getAnalysisIfAvailable<DominanceFrontier> ();
		if (DF)
			DF->runOnFunction(*F);
	}
	return true;
}
